import numpy as np

from .module import Module
from autograd import Tensor
from autograd.nn.parameter import Parameter
from .. import functional as F

class Conv2d(Module):
    r""" 对F.conv2d的包装 
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: tuple,
        stride: int = 1,
        padding: int = 0,
        dilation: int = 1,
        groups: int = 1,
        bias: bool = True,
        ):
        super(Conv2d, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.groups = groups
        self.weight = Parameter(Tensor.uniform(size=[out_channels, in_channels, kernel_size,kernel_size]))
        if bias:
            self.bias = Parameter(Tensor.uniform(size=out_channels))
        else:
            #self.register_parameter('bias', None)
            self.bias = None


    def forward(self, input):
        return F.conv2d(input, self.weight, self.bias, self.stride,
                        self.padding, self.dilation, self.groups)
        


